from flask import Flask, render_template, request, send_file
import torch
from transformers import BertForSequenceClassification, BertTokenizer
import pandas as pd
import io

app = Flask(__name__)

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载预训练的模型和tokenizer
model = BertForSequenceClassification.from_pretrained("./bert/sentiment_best_model_loss")
model.to(device)
tokenizer = BertTokenizer.from_pretrained("./bert/sentiment_best_model_loss")

# 定义预测函数
def predict_sentiment(text):
    encoding = tokenizer(text, return_tensors='pt', truncation=True, padding=True)
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    with torch.no_grad():
        output = model(input_ids, attention_mask=attention_mask)
        logits = output.logits
        predicted_label = torch.argmax(logits, dim=1).item()

    label_map = {0: 'positive', 1: 'negative', 2: 'neutral'}
    return label_map[predicted_label]

@app.route('/', methods=['GET', 'POST'])
def index():
    result = None
    if request.method == 'POST':
        if 'text' in request.form:
            text = request.form['text']
            result = predict_sentiment(text)
        elif 'file' in request.files:
            file = request.files['file']
            if file and file.filename.endswith('.csv'):
                df = pd.read_csv(file)
                if 'text' in df.columns:
                    df['sentiment'] = df['text'].apply(predict_sentiment)
                    output = io.BytesIO()
                    df.to_csv(output, index=False, encoding='utf-8-sig')
                    output.seek(0)
                    return send_file(output,
                                     mimetype='text/csv',
                                     as_attachment=True,
                                     download_name='sentiment_analysis_results.csv')
    return render_template('index.html', result=result)

if __name__ == '__main__':
    app.run(debug=True)